# coding:utf-8
import time

from keras.models import load_model
import keras
import os
import numpy as np
from scipy import misc

ROWS = 297
COLS = 396
CHANNELS = 3
mean = [0.485, 0.456, 0.406]


def main(model_path, testdir):
    model: keras.models.Model = load_model(model_path)
    imgs = os.listdir(testdir)
    data = np.ndarray((len(imgs), ROWS, COLS, CHANNELS), dtype=np.uint8)
    for i, imgpath in enumerate(os.listdir(testdir)):
        image = os.path.join(testdir, imgpath)
        img = misc.imread(image, mode="L")
        # 按需裁定尺寸
        img = misc.imresize(img, (ROWS, COLS))
        img = np.expand_dims(img, axis=2)
        data[i] = img

    print('start:', time.strftime('%Y/%m/%d %H:%M:%S', time.localtime(time.time())))

    predictions = model.predict(data)
    print('end:  ', time.strftime('%Y/%m/%d %H:%M:%S', time.localtime(time.time())))

    predictionIndexs = np.argmax(predictions, axis=1)
    print(predictionIndexs)
    results = []
    maxpredict = []
    for i in range(len(predictions)):
        predict = predictions[i][predictionIndexs[i]]
        if predict < tolerance:
            results.append('NONE')
        else:
            results.append(candidate[predictionIndexs[i]])
        maxpredict.append(predict)

    for img, i, j in zip(os.listdir(testdir), maxpredict, results):
        print(f'{img}\n{i}    :{j}', end='\n\n')


####################################
def _get_data_list(img_root=r"E:\Data\ZResource\Datasets\animals\animals"):
    dirs = os.listdir(img_root)
    dirs.sort()
    images = []
    labels = []
    for i, cdir in enumerate(dirs):
        cdir = os.path.join(img_root, cdir)
        for img in os.listdir(cdir):
            imp = os.path.join(cdir, img)
            images.append(imp)
            labels.append(i)
    return images, labels
def test():
    testdir=""
    from keras.models import load_model
    from util import plot
    model = load_model(output_model_file)
    imgs = os.listdir(testdir)
    data = np.ndarray((len(imgs), HEIGHT, WIDTH, nclasses), dtype=np.uint8)
    images, labels = _get_data_list(testdir)
    preds=[]
    reals=[]
    for (imp, label) in zip(images, labels):
        img = misc.imread(imp)
        img = np.expand_dims(img)
        pred: np.ndarray = model.predict(img)
        pred = np.ravel(pred)[0]
        preds.append(pred)
        reals.append(label)
        print(f"{imp} 预测结果:\n"
              f"预测: {pred}  -- 真实: {label}")

if __name__ == "__main__":
    model_path = r'F:\Resources\model\ws3_model.h5'
    testdir = r'F:\bigphoto\test_images'
    # candidate = ['ONE', 'TWO', 'THREE', 'FOUR', 'FIVE', 'OTHER']
    candidate = ['ONE', 'TWO', 'THREE', 'FOUR', 'FIVE', 'SIX', 'SEVEN', 'EIGHT']
    tolerance = 0.5
    main(model_path, testdir)
